# Copyright (c) 2021 Presto Labs Pte. Ltd.
# Author: hjkim
# pylint: disable=wrong-import-position,too-few-public-methods
# pylint: disable=unnecessary-comprehension,too-many-locals,broad-except
# pylint: disable=no-name-in-module
import os
from os.path import join as pjoin
import time
import json
import ujson
import datetime
import tempfile
import zipfile
import pandas
import copy

import fire

import python.coin_deploy.lm_agg2.kline_util as kline_util
from python.coin_deploy.lm_agg2.from_coin1_generator import is_focus_bnb_hack
from google.protobuf.json_format import MessageToDict

from xunkemgmt_client.tool.slack_noti import send_to_slack

os.environ["TZ"] = "GMT"
time.tzset()

intv_to_seconds = {
  "1m": 60,
  "5m": 300,
  "1h": 3600,
  "1d": 3600 * 24
}

def get_kline_request_dict_and_vu_rank(config_path):
  ret_dict = {}
  ret_rank = None
  confobj = ujson.load(open(config_path, 'r'))
  if 'products' not in confobj:
    return None
  for product_name, prodconf in confobj['products'].items():
    if is_focus_bnb_hack(product_name, prodconf):
      continue
    else:
      if 'profile_name' not in prodconf:
        continue
      profile_name = prodconf['profile_name']
      zipobj = zipfile.ZipFile(os.path.join("profile", f"{profile_name}.zip"))
      assert "feature_info.json" in [f.filename for f in zipobj.filelist]
      repr_dict = json.load(zipobj.open('feature_info.json'))
      zipobj.close()
      assert "features" in repr_dict
      for config_obj in [feature['config_obj'] for feature in repr_dict['features'] if 'kline_request_dict' in feature['config_obj']]:
        for key, val in config_obj['kline_request_dict'].items():
          ret_dict[key] = val if key not in ret_dict else max(ret_dict[key], val)
      if "focus_symbols_vu_rank" in repr_dict:
        ret_rank = str(repr_dict["focus_symbols_vu_rank"])

  return ret_dict, ret_rank

def get_vu_require_quote_matching(config_path):
  ret = False
  confobj = ujson.load(open(config_path, 'r'))
  for product_name, prodconf in confobj['products'].items():
    if is_focus_bnb_hack(product_name, prodconf):
      continue
    else:
      if 'profile_name' not in prodconf:
        continue
      profile_name = prodconf['profile_name']
      zipobj = zipfile.ZipFile(os.path.join("profile", f"{profile_name}.zip"))
      assert "feature_info.json" in [f.filename for f in zipobj.filelist]
      repr_dict = json.load(zipobj.open('feature_info.json'))
      zipobj.close()
      if "focus_vu_require_quote_matching" in repr_dict:
        ret = bool(repr_dict["focus_vu_require_quote_matching"])
  return ret

def get_symbol_universe(config_path, volume_info=None):
  require_quote_matching = get_vu_require_quote_matching(config_path)
  volume_lists = None
  if volume_info is not None:
    volume_lists = dict()
    for exch_volume_info in volume_info['volume_info_by_exchs']:
      tmp_list = sorted([(row['volume'], row['symbol']) for row in exch_volume_info['volume_infos']], reverse=True)
      volume_lists[(exch_volume_info['exchange_name'], exch_volume_info['market_type'])] = [symbol for _, symbol in tmp_list]

  ret = {}
  focus_exchange = None
  ref_exchanges = []
  ref_pns = []
  confobj = ujson.load(open(config_path, 'r'))
  if 'products' not in confobj:
    return None
  for product_name, prodconf in confobj['products'].items():
    if is_focus_bnb_hack(product_name, prodconf):
      continue
    else:
      if 'profile_name' not in prodconf:
        continue
      profile_name = prodconf['profile_name']
      zipobj = zipfile.ZipFile(os.path.join("profile", f"{profile_name}.zip"))
      assert "feature_info.json" in [f.filename for f in zipobj.filelist]
      repr_dict = json.load(zipobj.open('feature_info.json'))
      zipobj.close()
      assert "samplers" in repr_dict
      for config_obj in [feature['config_obj'] for feature in repr_dict['samplers']]:
        for symbol in config_obj['symbols']:
          if focus_exchange is None:
            focus_exchange = symbol['focus_exchange']
          if len(ref_exchanges) == 0:
            for i in [''] + list(range(2, 10)):
              if f'exchange{i}' not in symbol:
                break
              ref_exchanges.append(symbol[f'exchange{i}'])
              ref_pns.append(symbol[f'product_name{i}'])
          else:
            for i in [''] + list(range(2, 10)):
              if f'exchange{i}' not in symbol:
                break
              if i != '':
                assert ref_exchanges[i] == symbol[f'exchange{i}']
                assert ref_pns[i] == symbol[f'product_name{i}']
          assert focus_exchange == symbol['focus_exchange']

          # test for vu:
          # ./pyrunner cc/appcoin2/strategy/lm/kline_dumper_live.py \
          # run --config_path python/coin_deploy/lm_agg2/focus_vu/lm_focus_binance_usdt_swap_vu_1.json \
          # --kline_path kline_vu.json
          if symbol['focus_product_name'].find(".*") >= 0:
            import cc.appcoin2.strategy.lm.kline_generator as klg
            focus_mt = 'Futures' if focus_exchange.find('Futures') >= 0 else 'Spot'
            ref_mts = ['Futures' if ref_exchange.find('Futures') >= 0 else 'Spot' for ref_exchange in ref_exchanges]
            focus_symbols_pi = klg.get_symbols_from_pi(focus_mt, focus_exchange.replace("Futures",""))
            ref_symbols_pis = [
                klg.get_symbols_from_pi(ref_mt, ref_exchange.replace("Futures",""))
                for ref_mt, ref_exchange in zip(ref_mts, ref_exchanges)]

            def get_base_quote_pn_map(symbol_patt, symbols_pi, norm_base=False, volume_list=None):
              ret_map = {}
              import re
              symbol_re = re.compile(symbol_patt)
              for symbol_pi in symbols_pi:
                if re.match(symbol_patt, symbol_pi):
                  base = symbol_re.search(symbol_pi).group(1)
                  quote = kline_util.get_quote(symbol_pi) if require_quote_matching else ""
                  if norm_base:
                    base = kline_util.get_norm_base(base)
                  key = f"{base}-{quote}"
                  if key in ret_map and volume_list:
                    if ret_map[key] in volume_list and symbol_pi in volume_list and volume_list.index(ret_map[key]) > volume_list.index(symbol_pi):
                      ret_map[key] = symbol_pi
                  else:
                    ret_map[key] = symbol_pi
              return ret_map
            focus_base_quote_pn_map = get_base_quote_pn_map(symbol['focus_product_name'], focus_symbols_pi, norm_base=False)
            ref_base_quote_pn_maps = [
                get_base_quote_pn_map(ref_pn, ref_symbols_pi, norm_base=True, volume_list=volume_lists[(ref_exchanges[i].replace("Futures", ""), ref_mts[i])] if volume_lists else None)
                for i, (ref_pn, ref_symbols_pi) in enumerate(zip(ref_pns, ref_symbols_pis))]
            for base, focus_product_name in focus_base_quote_pn_map.items():
              ret[focus_product_name] = [
                  ref_base_quote_pn_map.get(kline_util.get_norm_base(base, preprocess=False), None)
                  for ref_base_quote_pn_map in ref_base_quote_pn_maps]
            continue

          if symbol['focus_product_name'] not in ret:
            ret[symbol['focus_product_name']] = symbol['product_name']
          assert ret[symbol['focus_product_name']] == symbol['product_name']

  return ret, focus_exchange, ref_exchanges

def is_enable_clear_all_1(config_path):
  keyword = "enable_clear_all"
  value = 1
  config_dict = json.loads(open(config_path).read())
  for product_key in config_dict["products"].keys():
    if keyword in config_dict["products"][product_key] and config_dict["products"][product_key][keyword] == value:
      return True
  return False

def run(now=None, kline_path=None, test=False, exchange="Binance",
        use_slack=False, tolerance=100, config_path=None, lm_strat_json_from_conf_gen=None):
  assert config_path is not None
  if is_enable_clear_all_1(config_path):
    with open(kline_path, "w", encoding="utf8") as fout:
      fout.write(json.dumps({"klines": []}, indent=2))
    return

  import cc.appcoin2.strategy.lm.kline_generator as klg
  kline_request_dict, vu_rank = get_kline_request_dict_and_vu_rank(config_path)

  el_st = time.time()
  if now is None:
    now = datetime.datetime.now()
  elif isinstance(now, (int, str)):
    now = datetime.datetime.strptime(str(now), "%Y%m%d%H%M")
  else:
    raise ValueError(now)

  sequential = False
  if "sequential" in kline_request_dict:
    sequential = bool(kline_request_dict.pop("sequential"))
  dump_ref = False
  if "dump_ref" in kline_request_dict:
    dump_ref = bool(kline_request_dict.pop("dump_ref"))
  print("Dump Ref: ", dump_ref)

  kline_request_dict = {
      key: kline_request_dict[key]
      for key in sorted(kline_request_dict.keys(), key=lambda x: intv_to_seconds[x])}

  print("kline period info:")
  start_dt_dict = {}
  end_dt_dict = {}
  end_dt_it = copy.deepcopy(now)
  for intv, num in kline_request_dict.items():
    start_dt_dict[intv] = end_dt_it - datetime.timedelta(seconds=intv_to_seconds[intv]*num)
    end_dt_dict[intv] = end_dt_it
    if sequential:
      end_dt_it = start_dt_dict[intv]

  kline_request_dict = {
      key: kline_request_dict[key]
      for key in sorted(kline_request_dict.keys(), key=lambda x: intv_to_seconds[x], reverse=True)}
  for intv, num in kline_request_dict.items():
    print(f"\tstart:\t{start_dt_dict[intv]}")
    print(f"\tend:\t{end_dt_dict[intv]} ({intv} * {num})")

  if dump_ref:
    print("dump_ref: on")

  volume_info = None
  if lm_strat_json_from_conf_gen is not None and vu_rank:
    if vu_rank == "all":
      vu_rank_ge, vu_rank_lt = 0, 10000
    else:
      vu_rank_ge, vu_rank_lt = list(map(int, vu_rank.split("-"))) if "-" in vu_rank else [0, int(vu_rank)]
    volume_info = json.loads(open(lm_strat_json_from_conf_gen).read()).get("volume_info", {})

  symbol_map, focus_exchange, ref_exchanges = get_symbol_universe(config_path, volume_info)
  focus_market_type = "Futures" if "Futures" in focus_exchange else "Spot"
  ref_market_types = ["Futures" if "Futures" in ref_exchange else "Spot" for ref_exchange in ref_exchanges]

  # filter symbols by vu logic
  if volume_info is not None:
    volume_info_by_exchs = volume_info.get("volume_info_by_exchs", [])
    symbol_map_whole = copy.deepcopy(symbol_map)
    symbol_map = {}
    if len(volume_info_by_exchs) > 0:
      volume_info_list = sorted([
          (volume_info_entry["volume"], volume_info_entry["symbol"])
          for volume_info_entry in volume_info_by_exchs[0].get("volume_infos", [])],
          reverse=True)
      cnt = 0
      print(pandas.DataFrame([{"focus": key, "ref": value} for key, value in symbol_map_whole.items()]).to_string())
      print(pandas.DataFrame(volume_info_list).to_string())
      for _, symbol in volume_info_list:
        if symbol_map_whole[symbol]:
          if dump_ref and (None in symbol_map_whole[symbol]):
            continue
          if vu_rank_ge <= cnt < vu_rank_lt or "BTC-" in symbol or "ETH-" in symbol:
            symbol_map[symbol] = symbol_map_whole[symbol]
          cnt += 1

  kline_info = {"klines": []}
  failed_info = {"failed": []}

  def convert_klinelist(klinelist):
    return MessageToDict(klinelist,
                        including_default_value_fields=True,
                        preserving_proto_field_name=True)

  for focus_symbol, ref_symbols in symbol_map.items():
    for intv, num in kline_request_dict.items():
      focus_klinelist, focus_failed = klg.query_kline_as_pb(
          now,
          f"{focus_market_type}:{focus_exchange.replace('Futures', '')}:{focus_symbol}",
          start_dt_dict[intv], end_dt_dict[intv], kline_period=intv, max_cnt=1000, tolerance=tolerance)
      focus_klinelist = convert_klinelist(focus_klinelist)
      failed_info["failed"] += focus_failed
      kline_info["klines"].append({
        "symbol": f"{focus_exchange}_{focus_symbol}",
        "interval_ns": int(intv_to_seconds[intv] * 1e9),
        "ohlc": focus_klinelist["klines"]})
      print(f"kline of {focus_symbol}(Focus) ({intv}) fetched")
      if dump_ref:
        for ref_symbol, ref_market_type, ref_exchange in zip(ref_symbols, ref_market_types, ref_exchanges):
          if ref_symbol is None:
            continue
          ref_klinelist, ref_failed = klg.query_kline_as_pb(
              now,
              f"{ref_market_type}:{ref_exchange.replace('Futures', '')}:{ref_symbol}",
              start_dt_dict[intv], end_dt_dict[intv], kline_period=intv, tolerance=tolerance)
          ref_klinelist = convert_klinelist(ref_klinelist)
          failed_info["failed"] += ref_failed
          kline_info["klines"].append({
            "symbol": f"{ref_exchange}_{ref_symbol}",
            "interval_ns": int(intv_to_seconds[intv] * 1e9),
            "ohlc": ref_klinelist["klines"]
          })
          print(f"kline of {ref_symbol}(Ref) ({intv}) fetched")
  dt_str = now.strftime("%Y%m%d%H")
  kline_path = kline_path or \
    pjoin(tempfile.mkdtemp(prefix=f"coin_midfreq_kline_{dt_str}", dir="/tmp"),
          "kline_info.json")
  print(f"dump kline to {kline_path}")
  with open(kline_path, "w", encoding="utf8") as fout:
    fout.write(json.dumps(kline_info, indent=2))
  elapsed = time.time() - el_st
  print(f"elapsed to get kline: {elapsed:.4f} sec")
  if use_slack:
    channel = "#k_hft_operation_log"
    if failed_info["failed"]:
      failed_df = pandas.DataFrame(failed_info["failed"])
      failed_cnt = failed_df['product'].value_counts().sort_values()
      send_to_slack('\n'.join([
                      f'Hostname: {os.environ.get("HOSTNAME", "")}',
                      f'kline_path: {kline_path}',
                      "",
                      "Failure counts",
                      failed_cnt.to_string(),
                      "",
                      "Failures",
                      failed_df.to_string()]),
                      channel,
                      'file',
                      now.strftime(f"Kline failures %Y%m%d_%H%M%S {focus_exchange.replace('Futures', '')}"))


if __name__ == "__main__":
  fire.Fire()
